#!/usr/bin/env python
#
# Scheduler watchdog for VOC simulations.
#
# Author:   Mike Murphy (mamurph@cs.clemson.edu)
# Revision: 27 September 2009
#
#   Copyright 2009 Clemson University
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
# This material is based upon work supported under a National Science
# Foundation Graduate Research Fellowship. Any opinions, findings,
# conclusions or recommendations expressed in this publication are those of
# the author(s) and do not necessarily reflect the views of the National
# Science Foundation.

'''
Watchdog component for VOC simulations. This component implements the policy
used to determine when the size of a VOC should be changed in response to
changing scheduler queue size.
'''

from results import VOCResult, WatchdogResult

DEFAULT_INTERVAL = 300


class WatchRecord(object):
   '''
   Watchdog record for storing information about each Virtual Organization
   known to the watchdog. Stored fields include a reference to the VOC
   owned by the VO, limit and target information, and a list of associated
   pilot jobs used to lease resources for the VOC.
   '''
   def __init__(self, voc, limit=-1, target=0):
      '''
      @param voc: Reference to the VOC owned by the VO
      @param limit: VOC size limit for the VO
      @type limit: int
      @param target: VOC size target for the VO
      @type target: int
      '''
      self.voc = voc
      self.limit = limit
      self.target = target
      self.pilot_jobs = []
   def get_running_count(self):
      '''
      @rtype: int
      @return: the number of pilot jobs submitted for the VO
      '''
      return len(self.pilot_jobs)
#


class Watchdog(object):
   '''
   Naive greedy watchdog implementation. This watchdog always tries to expand
   the size of a VOC whenever the number of jobs waiting in the watch queue
   exceeds the number of pilot jobs running to lease resources for the VOC,
   provided that the current number of pilot jobs does not exceed a specified
   limit level. Whenever the number of jobs waiting in the watch queue falls
   below the size of the VOC, the watchdog will terminate pilot jobs. The
   minimum size of a VOC (the target level) may be specified so that a VOC
   never becomes empty.

   This implementation supports tracking multiple VOs and thus is usable for
   both transparent and participating VOCs. However, only one scheduler queue
   may be watched, and only one scheduler queue may be used as a target for
   sending pilot jobs (both queues may be the same).
   '''
   def __init__(self, sim, name, w_sched, q_sched, \
                interval=DEFAULT_INTERVAL, default_target=0, \
                default_limit=-1, **kwds):
      '''
      @param sim: Reference to the simulator object
      @param name: Name of the watchdog
      @type name: str
      @param w_sched: Reference to the watch scheduler: this scheduler
                      is the one whose queue the watchdog monitors
      @param q_sched: Reference to the queue scheduler: this scheduler
                      is the one to which pilot jobs will be submitted
                      in order to lease additional resources to expand
                      the size of the VOC
      @param interval: Watchdog interval in time units
      @type interval: number
      @param default_target: Default target level to use when the VO input
                             record does not specify a target level
      @type default_target: int
      @param default_limit: Default limit level to use when the VO input
                            record does not specify a limit level
      @type default_limit: int
      @param kwds: Extra data for the watchdog
      '''
      self.sim = sim
      self.name = name
      self.w_sched = w_sched
      self.q_sched = q_sched
      self.interval = interval
      self.default_target = default_target
      self.default_limit = default_limit
      self.extra_data = kwds
      self.vo_map = {}            # name: watch record
      self.count_above_target = 0
      self.pending_event = None
      self.start_time = sim.get_time()
      self.last_watch_run = self.start_time
      self.w_sched.set_watchdog_callback(self.check_watch)
      self.check_watch()
   def debug(self, message):
      '''
      Transmits a debugging message to the simulator debug facility

      @param message: Message to be transmitted
      '''
      self.sim.debug('Watchdog<%s>: %s' % (self.name, message))
   def event_pending(self):
      '''
      @rtype: bool
      @return: True iff a watchdog event is pending
      '''
      wait = False
      if self.pending_event is not None:
         wait = not self.pending_event.activated
      return wait
   def add_watchdog_event(self, time):
      '''
      Adds a watchdog event at the specified time. Only one event may be
      pending at any time: attempting to schedule another event will result
      in an Exception.

      @param time: Time at which event should occur
      @type time: number
      '''
      if self.event_pending():
         raise Exception('Watchdog event already pending')
      self.pending_event = self.sim.new_periodic_event(time, self.watch_trigger)
   def check_watch(self):
      '''
      Adds a watchdog event if necessary. An event is necessary if no event
      is already pending, and either jobs are waiting in the watch scheduler
      queue or more VOC nodes are running than required by the target levels.

      @rtype: bool
      @return: True iff a watchdog event is added to the simulation
      '''
      result = False
      if not self.event_pending():
         need = (self.w_sched.get_waiting_count() > 0) or \
                (self.count_above_target > 0)
         if need:
	    cur_time = self.sim.get_time()
	    adjusted = self.last_watch_run - self.start_time
	    offset = self.interval - (adjusted % self.interval)
	    when = cur_time + offset
	    self.add_watchdog_event(when)
	    result = True
      return result
   def set_policy(self, wr, vo_record=None):
      '''
      Adjusts the policy contained in a particular WatchRecord instance by
      adjusting the target and limit levels as contained in the vo_record
      (dictionary-like object). If the vo_record is None, then this method
      does nothing. Any policy changes made by this method will not have
      any effect until the next watchdog invocation.

      @param wr: WatchRecord to adjust
      @param vo_record: VO record to use to make the adjustment
      '''
      if vo_record is not None:
         if vo_record.has_key('target'):
            wr.target = vo_record['target']
         if vo_record.has_key('limit'):
            wr.limit = vo_record['limit']
   def add_voc(self, voc, vo_record=None):
      '''
      Adds a VOC to the watchdog for management, or changes the policy of an
      existing VOC.

      @param voc: Reference to the VOC to add
      @param vo_record: Instance of L{vo_format_reader.VO}, or another
                        dictionary-like object, if available. This record
                        may contain fields named "limit" and "target",
                        which will be used to set or change the limit and
                        target levels for the corresponding VO.
      '''
      vo = voc.vo
      if vo in self.vo_map:
         rec = self.vo_map[vo]
         self.set_policy(rec, vo_record)
      else:
         rec = WatchRecord(voc, limit=self.default_limit, \
                           target=self.default_target)
         self.vo_map[vo] = rec
         self.set_policy(rec, vo_record)
   def start_voc_node(self, vo, count=1):
      '''
      Starts one or more VOC nodes for a VO managed by the watchdog, using the
      queue scheduler as the pilot job submission target.

      @param vo: Name of the VO
      @type vo: str
      @param count: Number of nodes to add to the VOC
      @type count: int
      '''
      record = self.vo_map[vo]
      voc = record.voc
      target = record.target
      record.pilot_jobs.extend(voc.grow(count, self.q_sched))
      self.debug('grew VOC for vo %s by %d' % (vo, count))
   def stop_voc_node(self, vo, count=1):
      '''
      Stops one or more VOC nodes for a VO managed by the watchdog. Note that
      this method will NOT stop VOC nodes that are presently busy, so the
      total number of VOC nodes that are stopped may be less than count.

      @param vo: Name of the VO
      @type vo: str
      @param count: Number of VOC nodes to stop
      @type count: int

      @return: a simulator results collection
      '''
      r = []
      record = self.vo_map[vo]
      target = record.target
      pilots = record.voc.shrink(count, kill_busy=False)
      for pilot in pilots:
         record.pilot_jobs.remove(pilot)
      self.debug('shrunk VOC for vo %s by %d' % (vo, count))
      return (r,)
   def watch_trigger(self, time, simobj):
      '''
      Periodic watchdog event handler that implements the naive greedy
      watchdog algorithm. This method is invoked as a simulator event.

      @param time: Current simulated time
      @type time: number
      @param simobj: Reference to the simulator object

      @return: simulator results collection
      '''
      rl = []
      self.count_above_target = 0
      for vo in self.vo_map.keys():
         record = self.vo_map[vo]
         qd = self.w_sched.get_total_queue_size(vo)
	 nr = record.get_running_count()
	 free = self.q_sched.get_free_count()
	 self.debug('vo: %s  qd: %d  nr: %d  free: %d' % (vo, qd, nr, free))
	 if (qd > nr) and (free > 0):
            if (record.limit < 0) or (nr < record.limit):
               count = min(free, (qd - nr))
               if record.limit > 0:
                  count = min(count, (record.limit - nr))
               self.start_voc_node(vo, count)
	 elif (record.target + qd) < nr:
	    surplus = nr - (record.target + qd)
            rl.append(self.stop_voc_node(vo, surplus))
         #
         above_target = record.get_running_count() - record.target
         if above_target > 0:
            self.count_above_target += above_target
      #
      self.last_watch_run = time
      self.check_watch()
      #
      rslt = None
      w_size = self.w_sched.get_total_queue_size()
      q_size = self.q_sched.get_total_queue_size()
      w_nametag = self.w_sched.name + '@' + self.w_sched.tag
      q_nametag = self.q_sched.name + '@' + self.q_sched.tag
      watch_res = WatchdogResult(time, self.name, w_nametag, q_nametag, \
	                         w_size, q_size)
      rl.append(watch_res)
      for vo in self.vo_map:
         rec = self.vo_map[vo]
	 count = rec.get_running_count()
         vo = rec.voc.vo
	 vo_res = VOCResult(time, vo, count, rec.target, rec.limit)
         rec.last_running = count
	 rl.append(vo_res)
      rslt = (rl,)
      return rslt
#


if __name__ == '__main__':
   from cluster_sim import Cluster, MachineFactory
   from condor_sim import Condor
   from job_model import Job
   from kernel import Simulator
   from storage import FileCache, FileStore
   from units import *
   from virtualization import VMHost, VMImage
   from voc_sim import VOC

   sim = Simulator()
   mf = MachineFactory(VMHost, vmm_name='KVM', sim=sim, cpu_oh=9.0)
   c1 = Cluster('metal', sim, factory=mf)
   c1.grow(10)
   sched1 = Condor(sim)
   c1.add_scheduler('Scheduler', sched1)
   sched1.add_targets(c1.machines)

   sched2 = Condor(sim)

   fs = FileStore('VMStore', sim, 1*u_Gbps)
   img = VMImage('vm.img', 5*u_GiB, sched2, 60)
   fs.add_file(img)

   fc = FileCache('VMCache', sim, 1*u_Gbps)
   sched1.add_resource('vmstore', fc)

   c2 = VOC('cluster2', sim, 'engage', 'vm.img', fs)
   wd = Watchdog(sim, 'watchdog', sched2, sched1)
   wd.add_voc(c2)

   for x in xrange(0,20):
      job = Job(20 + (x % 4), name='job%d' % x, vo='engage', prob_err=0.2)
      sched2.add_submit_event(10*x, job)
   #

   sim.run()
#
